#include <string.h>

#include "tdp2_allocator.hpp"
#include "tdp2_memzone.hpp"

namespace TDP
{
	namespace MemoryManagement
	{
		void *MemoryAllocator::realloc(void *opaque, void *ptr, LargeUInt sz, LargeUInt alignment)
		{
			MemoryAllocator *ma = static_cast<MemoryAllocator*>(opaque);

			if(ptr == NULL)
			{
				// Allocate
				if(sz == 0)
					return NULL;
				return ma->Alloc(sz, alignment);
			}

			if(sz == 0)
			{
				// Free
				ma->Free(ptr);
				return NULL;
			}

			// Attempt to resize
			LargeUInt usedSpace;
			if(ma->Resize(ptr, sz, usedSpace))
				return ptr;

			// Failed, try to relocate to a new block
			void *newMem = ma->Alloc(sz, alignment);
			if(newMem == NULL)
				return NULL;	// Failed

			// Copy the old data
			memcpy(newMem, ptr, usedSpace);
			// Free the old data
			ma->Free(ptr);

			// Done
			return newMem;
		}

		//*********************************************************
		// Zone allocator
		ZoneAllocator::ZoneAllocator(void *memBlock, LargeUInt memBlockSize, LargeUInt blockAlignment)
		{
			_zoneBytes = static_cast<UInt8*>(memBlock);

			LargeUInt misalignment = PadForAddress(_zoneBytes, RDX_ALIGNOF(ZoneBlock));

			_zoneBytes += misalignment;
			memBlockSize -= misalignment;
			_endOfBlock = _zoneBytes + memBlockSize;
			_lastBlock = reinterpret_cast<ZoneBlock*>(_zoneBytes);
			_blockAlignment = blockAlignment;
			_totalBytes = memBlockSize;

			_lastBlock->blockSize = memBlockSize;
			_lastBlock->prevBlockOffset = 0;
			_lastBlock->usedSpace = sizeof(ZoneBlock);
		}

		bool ZoneAllocator::Resize(void *ptr, LargeUInt sz, LargeUInt &outUsedSpace)
		{
			LargeUInt headerSize;
			ZoneBlock *block = PtrToBlock(ptr, headerSize);

			// See if the new requested size fits in the current block
			if(sz <= block->blockSize - headerSize)
			{
				// It does
				block->usedSpace = headerSize + sz;
				return true;
			}
			// It doesn't
			outUsedSpace = block->usedSpace - headerSize;
			return false;
		}

		// Tries to split a block into two blocks, the second containing the new data
		// If successful, returns the block
		void *ZoneAllocator::AllocInBlock(ZoneAllocator::ZoneBlock *block, LargeUInt sz, LargeUInt align)
		{
			if(block->blockSize - block->usedSpace <= sz)
				return NULL;	// Don't even try, this is too small

			UInt8 *blockBase = reinterpret_cast<UInt8*>(block);
			LargeUInt headerOffset = block->usedSpace + PadForAddress(blockBase + block->usedSpace, _blockAlignment);
			LargeUInt dataOffset = headerOffset + sizeof(ZoneBlock) + 1;
			LargeUInt alignPad = PadForAddress(blockBase + dataOffset, align);
			dataOffset += alignPad;

			if(dataOffset > block->blockSize)
				return NULL;	// Data would start outside of the block
			if(block->blockSize - dataOffset < sz)
				return NULL;	// Not enough space for the new data

			// Cleared to allocate it
			ZoneBlock *newBlock = reinterpret_cast<ZoneBlock*>(blockBase + headerOffset);
			newBlock->prevBlockOffset = headerOffset;
			newBlock->blockSize = block->blockSize - headerOffset;
			newBlock->usedSpace = dataOffset + sz;

			// Shrink the old block
			block->blockSize = headerOffset;

			// Set the padding scrollback byte
			blockBase[dataOffset-1] = static_cast<UInt8>(alignPad + 1);

			// Make this the last block if we're splitting the last one
			if(block == _lastBlock)
				_lastBlock = newBlock;

			return blockBase + dataOffset;
		}

		ZoneAllocator::ZoneBlock *ZoneAllocator::PtrToBlock(void *ptr, LargeUInt &outHeaderSize)
		{
			// Find the alignment byte
			UInt8 *ptrBase = static_cast<UInt8*>(ptr);
			UInt8 padSize = ptrBase[-1];

			// Locate the header and return it
			outHeaderSize = static_cast<LargeUInt>(padSize) + sizeof(ZoneBlock);

			return reinterpret_cast<ZoneBlock*>(ptrBase - static_cast<LargeUInt>(padSize) - sizeof(ZoneBlock));
		}

		void *ZoneAllocator::Alloc(LargeUInt sz, LargeUInt align)
		{
			if(sz >= _totalBytes)
				return NULL;

			// Start at the last block
			ZoneBlock *scanBlock = _lastBlock;
			ZoneBlock *lastBlock = _lastBlock;
			while(true)
			{
				void *mem = AllocInBlock(scanBlock, sz, align);
				if(mem != NULL)
					return mem;

				if(scanBlock == lastBlock)
					scanBlock = reinterpret_cast<ZoneBlock*>(_zoneBytes);	// Go back to the first one
				else
					scanBlock = reinterpret_cast<ZoneBlock*>(reinterpret_cast<UInt8*>(scanBlock) + scanBlock->blockSize);

				if(scanBlock == lastBlock)
					return NULL;	// Looped all the way around and found nothing, fail out
			}
		}

		void ZoneAllocator::Free(void *ptr)
		{
			// Merge this into the previous block
			ZoneBlock *block = PtrToBlock(ptr);
			ZoneBlock *prevBlock = reinterpret_cast<ZoneBlock*>(reinterpret_cast<UInt8*>(block) - block->prevBlockOffset);
			prevBlock->blockSize += block->blockSize;

			// Move the last block back if this was the last one
			if(block == _lastBlock)
				_lastBlock = prevBlock;
		}


		//********************************************************************************************
		// Free list allocator
		FreeListAllocator::FreeListAllocator(void *space, LargeUInt size, LargeUInt unitSize, LargeUInt unitAlign)
		{
			_firstFree = NULL;

			if(unitSize < unitAlign)
				unitSize = unitAlign;

			UInt8 *spaceBase = static_cast<UInt8*>(space);
			LargeUInt padBase = PadForAddress(spaceBase, RDX_ALIGNOF(FreeListLink));

			if(padBase >= size)
				return;

			spaceBase += padBase;
			size -= padBase;

			// Determine the most amount of padding necessary for the unit alignment
			LargeUInt wastedSpace = 0;
			if(unitAlign > RDX_ALIGNOF(FreeListLink))
				wastedSpace = unitAlign - RDX_ALIGNOF(FreeListLink);

			if(wastedSpace >= size)
				return;

			LargeUInt spaceEstimator = size - wastedSpace;

			// Determine how many units are possible
			LargeUInt numUnits = (size - wastedSpace) / (sizeof(FreeListLink) + unitSize);

			// Allocate actual space
			FreeListLink *links = reinterpret_cast<FreeListLink*>(spaceBase);

			UInt8 *unitBase = spaceBase + sizeof(FreeListLink) * numUnits;
			unitBase += PadForAddress(unitBase, unitAlign);

			_units = unitBase;
			_unitAlign = unitAlign;
			_unitSize = unitSize;

			// Initialize the free list
			_firstFree = links;
			for(LargeUInt i=0;i<numUnits-1;i++)
				links[i].next = links + i + 1;
			links[numUnits-1].next = NULL;

			_links = links;
		}

		// These blocks are unresizable so if the requested amount is within the block size, it's OK
		bool FreeListAllocator::Resize(void *ptr, LargeUInt size, LargeUInt &outUsedSpace)
		{
			outUsedSpace = _unitSize;
			return (size <= _unitSize);
		}

		void *FreeListAllocator::Alloc(LargeUInt sz, LargeUInt align)
		{
			FreeListLink *link = _firstFree;
			if(link == NULL)
				return NULL;

			// Unlink from the free list
			_firstFree = link->next;

			return _units + _unitSize * (link - _links);
		}

		void FreeListAllocator::Free(void *ptr)
		{
			FreeListLink *link = _links + ((static_cast<UInt8*>(ptr) - _units) / _unitSize);

			// Push on the free list
			link->next = _firstFree;
			_firstFree = link;
		}

		//****************************************************************************************
		// Multi-allocator
		MultiAllocator::MultiAllocator(const Region *regions, LargeUInt numRegions)
		{
			_regions = regions;
			_numRegions = numRegions;
		}

		bool MultiAllocator::Resize(void *ptr, LargeUInt size, LargeUInt &outUsedSpace)
		{
			LargeUInt numRegions = _numRegions;
			for(LargeUInt i=0;i<numRegions;i++)
			{
				if(_regions[i].ContainsPointer(ptr))
					return _regions[i].opaque->Resize(ptr, size, outUsedSpace);
			}
			outUsedSpace = 0;
			return false;
		}

		void MultiAllocator::Free(void *ptr)
		{
			LargeUInt numRegions = _numRegions;
			for(LargeUInt i=0;i<numRegions;i++)
			{
				if(_regions[i].ContainsPointer(ptr))
				{
					_regions[i].opaque->Free(ptr);
					return;
				}
			}
		}

		void *MultiAllocator::Alloc(LargeUInt sz, LargeUInt align)
		{
			LargeUInt numRegions = _numRegions;
			for(LargeUInt i=0;i<numRegions;i++)
			{
				void *mem = _regions[i].opaque->Alloc(sz, align);
				if(mem != NULL)
					return mem;
			}
			return NULL;
		}
	}
}
